import {
  TCompareFunc,
} from '../utils/tools';
import {
  LESS_THAN,
  BIGGER_THAN,
  BALANCE_FACTOR,
} from '../utils/constant';
import {
  BinarySearchTree,
} from './bstree';
import { TBinaryTreeNode, BinaryTreeNode } from './base';

export class AVLTree extends BinarySearchTree {
  public constructor(elements: any[], funcCompare?: TCompareFunc) {
    super(elements, funcCompare);
  }

  public getBalanceFactor(target: TBinaryTreeNode | any): BALANCE_FACTOR | undefined {
    const node = target instanceof BinaryTreeNode ? target : this.searchNode(target, this.root)?.node;

    if (node) {
      const heightDiff = this.getNodeHeight(node.left) - this.getNodeHeight(node.right);

      switch(heightDiff) {
        case -2:
          return BALANCE_FACTOR.UNBALANCED_RIGHT;
        case -1:
          return BALANCE_FACTOR.SIGHTLY_RIGHT;
        case 1:
          return BALANCE_FACTOR.SIGHTLY_LEFT;
        case 2:
          return BALANCE_FACTOR.UNBALANCED_LEFT;
        default:
          return BALANCE_FACTOR.BALANCED;
      }
    } else {
      return undefined;
    }
  }

  protected rotateLL(node: TBinaryTreeNode): TBinaryTreeNode {
    let temp = node?.left;
    if (node && temp) {
      node.left = temp.right;
      temp.right = node;
      return temp;
    } else {
      return undefined;
    }    
  }

  protected rotateRR(node: TBinaryTreeNode): TBinaryTreeNode {
    let temp = node?.right;
    if (node && temp) {
      node.right = temp.left;
      temp.left = node;
      return temp;
    } else {
      return undefined;
    }    
  }

  protected rotateLR(node: TBinaryTreeNode): TBinaryTreeNode {
    if (node) {
      node.left = this.rotateRR(node.left); 
      return this.rotateLL(node);
    } else {
      return undefined;
    }    
  }

  protected rotateRL(node: TBinaryTreeNode): TBinaryTreeNode {
    if (node) {
      node.right = this.rotateLL(node.right);
      return this.rotateRR(node);
    } else {
      return undefined;
    }    
  }

  protected insertNode(node: TBinaryTreeNode, key: any) {
    let newNode = super.insertNode(node, key);

    if (newNode) {
      const balanceFactor = this.getBalanceFactor(newNode);
      
      if (balanceFactor === BALANCE_FACTOR.UNBALANCED_LEFT) {
        const cmpresult = this.funcCompare(key, newNode.left?.key);
        if (cmpresult === LESS_THAN) {
          newNode = this.rotateLL(newNode);
        } else {
          newNode = this.rotateLR(newNode);
        }
      } else if (balanceFactor === BALANCE_FACTOR.UNBALANCED_RIGHT) {
        const cmpresult = this.funcCompare(key, newNode.right?.key);
        if (cmpresult === BIGGER_THAN) {
          newNode = this.rotateRR(newNode);
        } else {
          newNode = this.rotateRL(newNode);
        }
      }
    }

    return newNode;
  }

  protected removeNode(key: any, node: TBinaryTreeNode) {
    let nodeReturn = super.removeNode(key, node);
    
    if (nodeReturn) {
      const balanceFactor = this.getBalanceFactor(nodeReturn);
  
      if (balanceFactor === BALANCE_FACTOR.UNBALANCED_LEFT) {
        const balanceLeft = this.getBalanceFactor(nodeReturn.left);

        if (balanceLeft === BALANCE_FACTOR.UNBALANCED_LEFT) {
          nodeReturn = this.rotateLL(nodeReturn);
        } else {
          nodeReturn = this.rotateLR(nodeReturn);
        }
      } else if(balanceFactor === BALANCE_FACTOR.UNBALANCED_RIGHT) {
        const balanceLeft = this.getBalanceFactor(nodeReturn.right);

        if (balanceLeft === BALANCE_FACTOR.UNBALANCED_LEFT) {
          nodeReturn = this.rotateRR(nodeReturn);
        } else {
          nodeReturn = this.rotateLR(nodeReturn);
        }
      }
    }

    return nodeReturn;
  }
}
